#include <stdlib.h>
#include <stdio.h>
#include "cuda_runtime.h"
#include "nccl.h"

#if 1
#define CUDACHECK(cmd) do {                         \
  cudaError_t err = cmd;                            \
  if (err != cudaSuccess) {                         \
    printf("Failed: Cuda error %s:%d '%s'\n",       \
        __FILE__,__LINE__,cudaGetErrorString(err)); \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)


#define NCCLCHECK(cmd) do {                         \
  ncclResult_t res = cmd;                           \
  if (res != ncclSuccess) {                         \
    printf("Failed, NCCL error %s:%d '%s'\n",       \
        __FILE__,__LINE__,ncclGetErrorString(res)); \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)
#endif
void show_h_A_value(int *h_A, int cnt)
{
    for (int i = 0; i < cnt; ++i) {
        printf("%d ", (int)h_A[i]);
    }
    printf("\n");
}
#define MIN_ARGC                        (1)
#define MAX_SUPPORT_DEV_NUM_COUNT       (8)
char sCmdGetGPUNumber[] = "nvidia-smi topo -m";
int aSupportDevNum[MAX_SUPPORT_DEV_NUM_COUNT] = {1, 2, 4, 8};

int main(int argc, char* argv[])
{
  ncclComm_t comms[4];

  //managing 4 devices
  int nDev = 1;
  int count = 8;
  int devs[4] = { 0, 1, 2, 3 };

  int* sendbuff_H = (int*)malloc(count * sizeof(int));
  int* recvbuff_H = (int*)malloc(count * sizeof(int));
  //allocating and initializing device buffers
  int** sendbuff = NULL;
  int** recvbuff = NULL;
  cudaStream_t* stream = NULL;

  printf("=====================start==============\n");
  printf("[%s +%d %s] ncclAllReduce API test example\n", __FILE__, __LINE__, __func__);

  if(1 == argc)
  {
    printf("[%s +%d %s] argc=%d !=%d so exit(0)\n", __FILE__, __LINE__, __func__, argc, MIN_ARGC);
    printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
    printf("./nccl_send_recv.elf 1 \n");
    printf("./nccl_send_recv.elf 2 \n");
    printf("./nccl_send_recv.elf 4 \n");
    system(sCmdGetGPUNumber);
    exit(0);
  }
  else
  {
    int i = 0;
    nDev = atoi(argv[1]);
    printf("[%s +%d %s] input nDev=%d\n", __FILE__, __LINE__, __func__, nDev);
    printf("[%s +%d %s] MAX_SUPPORT_DEV_NUM_COUNT=%d\n", __FILE__, __LINE__, __func__, MAX_SUPPORT_DEV_NUM_COUNT);
    for(i = 0; i < MAX_SUPPORT_DEV_NUM_COUNT; i++)
    {
      if(nDev == aSupportDevNum[i])
      {
        break;
      }
    }
    if(MAX_SUPPORT_DEV_NUM_COUNT == i)
    {
      printf("[%s +%d %s] nDev=%d is not in aSupportDevNum[] so exit(0)\n", __FILE__, __LINE__, __func__, nDev);
      printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
      printf("./nccl_send_recv.elf 1 \n");
      printf("./nccl_send_recv.elf 2 \n");
      printf("./nccl_send_recv.elf 4 \n");
      printf("./nccl_send_recv.elf 8 \n");
      system(sCmdGetGPUNumber);
      exit(0);
    }
  }
  sendbuff = (int**)malloc(nDev * sizeof(int*));
  recvbuff = (int**)malloc(nDev * sizeof(int*));
  stream = (cudaStream_t*)malloc(sizeof(cudaStream_t)*nDev);
  
  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaMalloc((void**)sendbuff + i, count * sizeof(int)));
    CUDACHECK(cudaMalloc((void**)recvbuff + i, count * sizeof(int)));

    for (int j = 0; j < count; ++j)
    {
      *(sendbuff_H + j) = ((j + 1)*100) + ((i+1) * 10);
      printf("[%s +%d %s] [%d] *(sendbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(sendbuff_H + j));
    }   
    printf("[%s +%d %s] copy sendbuff_H to sendbuff by cudaMemcpy\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemcpy(sendbuff[i], sendbuff_H, count * sizeof(int), cudaMemcpyHostToDevice));

    printf("[%s +%d %s] set 0 to recvbuff by cudaMemset\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemset(recvbuff[i], 0, count * sizeof(int)));
    CUDACHECK(cudaStreamCreate(stream+i));
    printf("[%s +%d %s] copy recvbuff to recvbuff_H by cudaMemcpy\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemcpy(recvbuff[i], recvbuff_H, count * sizeof(int), cudaMemcpyHostToDevice));
    for (int j = 0; j < count; ++j)
    {
      printf("[%s +%d %s] [%d] *(recvbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(recvbuff_H + j));
    }   
  }

  printf("[%s +%d %s] call ncclCommInitAll\n", __FILE__, __LINE__, __func__);
  //initializing NCCL
  NCCLCHECK(ncclCommInitAll(comms, nDev, devs));

  //calling NCCL communication API. Group API is required when using
   //multiple devices per thread
  NCCLCHECK(ncclGroupStart());
  for (int i = 0; i < nDev; ++i)
  {
    printf("[%s +%d %s] call cudaSetDevice and ncclAllReduce\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaSetDevice(i));
  printf("[%s +%d %s] ncclAllReduce buffer index=%d\n", __FILE__, __LINE__, __func__, i);
    NCCLCHECK(ncclAllReduce((const void*)sendbuff[i], (void*)recvbuff[i], count, ncclInt, ncclSum, comms[i], stream[i]));
  }
  NCCLCHECK(ncclGroupEnd());

  //synchronizing on CUDA streams to wait for completion of NCCL operation
  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaStreamSynchronize(stream[i]));
  }

  for (int i = 0; i < nDev; ++i)
  {
    printf("[%s +%d %s] copy recvbuff[%d] to recvbuff_H by cudaMemcpy\n", __FILE__, __LINE__, __func__, i);
    CUDACHECK(cudaMemcpy(recvbuff_H, recvbuff[i], count * sizeof(int), cudaMemcpyDeviceToHost));
    for (int j = 0; j < count; ++j)
    {
      printf("[%s +%d %s] [%d] *(recvbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(recvbuff_H + j));
    }
  }

  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    //free device buffers
    CUDACHECK(cudaFree(sendbuff[i]));
    CUDACHECK(cudaFree(recvbuff[i]));
    //finalizing NCCL
    NCCLCHECK(ncclCommDestroy(comms[i]));
  }

  if(NULL != sendbuff_H)
  {
    free(sendbuff_H);
    sendbuff_H = NULL;
  }
  if(NULL != recvbuff_H)
  {
    free(recvbuff_H);
    recvbuff_H = NULL;
  }

  printf("=====================end==============\n");
  return 0;
}
